import torch
import Net
from Net.SR_MGANet import Gen_Guided_UNet
from Net.multiscaleloss import multiscaleEPE
import DataLoader.DataUtils as DataUtils
import DataLoader.SimuDataLoader as SimuDataLoader

import numpy as np
import matplotlib.pyplot as plt
import imp
import sys
import skimage.measure as Measure
import os
import sys
sys.path.append("./")
import clear_outputfiles as ClearFiles
from Net.EdgeColorLoss import EdgeColorLoss
import argparse

class TrainInBackGround(object):
    LR = 0.0001
    
    SCALE = 4
    
    ITER_NUM = 60 # Number of Iterations (Net trained with an epoch)
    EPOCH_SIZE = 880 # Number of batches in an epoch
    VAL_SET_SIZE = 350 # Number of batches in Valid Size
    BATCH_SIZE = 6 # Number of videos in a mini-batch
    LSTM_FRAME_NUM = 7
    IMG_SIZE = (200, 400)
    CUDA_DEVICE_IDX = 0
    
    EDGE_WEIGHT = 1
    
    
    COLOR_LOSS_MILESTONES = [0, ITER_NUM//2, ITER_NUM//4*3]
    COLOR_LOSS_VARIABLE_WEIGHT = [0.1, 0.2, 0.4]#[0.05, 0.1, 0.2]
    
    
    EDGE_LOSS_MILESTONES = [0, ITER_NUM//5, ITER_NUM//2]
    EDGE_LOSS_VARIABLE_WEIGHT = [0.1, 0.15, 0.2]#[0.05, 0.1, 0.2]
    
    
    def __init__(self, save_every=20):
        self.SAVE_EVERY = save_every
        
        self.LOGPATH = "./logs/"+DataUtils.get_time_stamp()+".log"

    def luckyprint(self, data, p=1, fp=None):
        if fp==None:
            fp = self.LOGPATH
        logfile = open(fp, "a")
        seed = np.random.random()
        if seed < p:
            if (logfile != None):
                logfile.write(DataUtils.get_time_stamp() + " " + str(data) + "\n")
            print(DataUtils.get_time_stamp() + " " + str(data))
        logfile.close()
        

    def loadNet(self):
        torch.cuda.set_device(self.CUDA_DEVICE_IDX)
        MGANet = Gen_Guided_UNet(input_size=[self.IMG_SIZE[0], 
                                             self.IMG_SIZE[1]], 
                                 batchNorm=False,
                                 scale_factor=self.SCALE).cuda()

        optimizer = torch.optim.Adam(MGANet.parameters(), lr=self.LR)
        self.luckyprint("Net Loaded")

        variableLR = torch.optim.lr_scheduler.MultiStepLR(optimizer, 
                                    milestones=[int(self.ITER_NUM/2), int(self.ITER_NUM*3/4)], gamma=0.1)

        return MGANet, optimizer, variableLR


    def valid(self, batch_index, net):
        net.eval()
        half_num = (self.LSTM_FRAME_NUM-1)//2 # frames number before or after this current frame
        
        x, y = SimuDataLoader.generateBatch_FromSavedFiles(
            self.EPOCH_SIZE*self.BATCH_SIZE + batch_index%self.VAL_SET_SIZE,
            BATCH_SIZE=1, imsize=self.IMG_SIZE,
            scale=self.SCALE)
        
        preds = np.zeros(y.shape)
        tensor_x = torch.from_numpy(x).cuda().float()
        tensor_y = torch.from_numpy(y).cuda().float()
            
        with torch.no_grad():
            prediction = net(tensor_x)
            x_us = prediction[5] # upsampled x for image compare
            preds = prediction[4].detach().cpu().numpy()[:]
            
        edge_residue = multiscaleEPE(
            prediction[:5], tensor_y[0], 
            L_model = EdgeColorLoss())[-1].numpy()
        
        frame_composition = np.mean(x_us[0], axis=0)
        DataUtils.show_bright_images([[x_us[0][half_num], x_us[0][0], frame_composition],
                                      [preds[0], edge_residue[0], y[0][0]]], 
                                     filename_prefix="valid_%d_"%batch_index, enhanced=False)
        self.luckyprint("Show pic")

        return (
            DataUtils.compare_video(np.clip(preds, 0, 1), 
                                    np.clip(y[0], 0, 1)), 
            DataUtils.compare_video(np.clip(x_us[:,half_num], 0, 1), 
                                    np.clip(y[0], 0, 1))
                )

    def train_step(self, batch_index, net, optimizer, thisColorWeight=0.1, thisEdgeWeight=0.1):
        net.train()

        this_train_loss = []
        half_num = (self.LSTM_FRAME_NUM-1)//2 # frames number before or after this current frame
        # half_num = 3 if 7 frmaes are to be passed to LSTM net
        
        x, y = SimuDataLoader.generateBatch_FromSavedFiles(batch_index, 
                                                           BATCH_SIZE=self.BATCH_SIZE, 
                                                           imsize=self.IMG_SIZE,
                                                           scale=self.SCALE)
            
        preds = np.zeros(y.shape)
        optimizer.zero_grad()
        tensor_x = torch.from_numpy(x).cuda().float()
        tensor_y = torch.from_numpy(y).cuda().float()
        
        prediction = net(tensor_x)
        loss, basic_loss, color_loss, edge_loss, edge_residue = multiscaleEPE(
            prediction[:5], tensor_y[:,0], 
            L_model = EdgeColorLoss(
                colorWeight=thisColorWeight,
                edgeWeight=thisEdgeWeight)
            )
        loss.backward()
        optimizer.step()
        
        this_train_loss.append([loss.item(), basic_loss.item(), color_loss.item(), edge_loss.item()])
        
        return np.mean(this_train_loss, axis=0)
    
    def generate_color_edge_weight(self, iter_idx):
        color_weight = 0
        for index in range(len(self.COLOR_LOSS_MILESTONES)):
            if iter_idx >= self.COLOR_LOSS_MILESTONES[index]:
                color_weight = self.COLOR_LOSS_VARIABLE_WEIGHT[index]
                
        edge_weight = 0
        for index in range(len(self.EDGE_LOSS_MILESTONES)):
            if iter_idx >= self.EDGE_LOSS_MILESTONES[index]:
                edge_weight = self.EDGE_LOSS_VARIABLE_WEIGHT[index]
                
        return color_weight, edge_weight
    
    def train(self):
        self.luckyprint("Current PID: " + str(os.getpid()))
        MGANet, optimizer, variableLR = self.loadNet()

        imp.reload(Net.multiscaleloss)  # reload module in case the file updates
        torch.cuda.empty_cache()
        train_loss = []
        for epoch_idx in range(self.ITER_NUM):
            this_color_weight, this_edge_weight = self.generate_color_edge_weight(epoch_idx)
            for i in range(self.EPOCH_SIZE):  
                this_epoch_train_loss = self.train_step(i, MGANet, optimizer, 
                                                        thisColorWeight=this_color_weight,
                                                        thisEdgeWeight=this_edge_weight)
                train_loss.append(this_epoch_train_loss)
                
                (psnr_pred, ssim_pred), (psnr_y, ssim_y) = (0, 0), (0, 0)
                
                this_LR = optimizer.param_groups[0]['lr']
                if i % self.SAVE_EVERY == 0 and i != 0:
                    (psnr_pred, ssim_pred), (psnr_y, ssim_y) = self.valid(i, MGANet)
                    
                    valstr = ('val_psnr: {:0.6f} / {:0.6f}, val_ssim: {:0.6f} / {:0.6f}').format(
                                psnr_pred, psnr_y, 
                                ssim_pred, ssim_y)
                    self.luckyprint(valstr)
                    
                if i % 10 == 0:
                    logstr = ('iter: [{:d}/{:d}], batch: [{:d}/{:d}], train_loss: {:0.6f}, basic_loss: {:0.6f}, color_loss: {:0.6f}, edge_loss: {:0.6f}, lr: {:.7f}, color_weight: {:.2f}, edge_weight: {:.2f}').format(
                        epoch_idx + 1, self.ITER_NUM,
                        (i + 1), self.EPOCH_SIZE,
                        train_loss[-1][0],
                        train_loss[-1][1],
                        train_loss[-1][2],
                        train_loss[-1][3],
                        this_LR, this_color_weight, this_edge_weight)
                    self.luckyprint(logstr)

            if epoch_idx % (self.ITER_NUM / 5) == (self.ITER_NUM / 5) - 1:
                torch.save(MGANet.state_dict(),"./pklmodels/train_iter_"+str(epoch_idx+1)+".pkl")

            variableLR.step()
            
            torch.cuda.empty_cache()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--clear', type=str, choices=('YES','NO'), default='NO') # clear previous output files/images
    parser.add_argument('--save_every', type=int, default=1) # save figure to file every x epoch
    args = parser.parse_args()
    
    if args.clear == 'YES':
        ClearFiles.clear_files()
        print("Clear Previous Output Files")
    
    if not os.path.exists("./pic/"):
        os.makedirs("./pic/")
    if not os.path.exists("./logs/"):
        os.makedirs("./logs/")
    if not os.path.exists("./pklmodels/"):
        os.makedirs("./pklmodels/")

    my_train = TrainInBackGround(save_every=args.save_every)

    file = open(my_train.LOGPATH, "w+")

    class Logger(object):
        def __init__(self, filename='default.log', stream=sys.stdout):
            self.terminal = stream
            self.log = open(filename, 'a')

        def write(self, message):
            self.terminal.write(message)
            self.log.write(message)

        def flush(self):
            pass


    sys.stdout = Logger(my_train.LOGPATH, sys.stdout)
    sys.stderr = Logger(my_train.LOGPATH, sys.stderr)

    my_train.train()
